from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

app = Flask(__name__)

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-zh")
model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-zh")
translation = pipeline('translation_zh_to_en', model=model, tokenizer=tokenizer)

# Define the translation endpoint
@app.route('/translate', methods=['GET'])
def translate():
    try:
        text_to_translate = request.args.get('text')
        result = translation(text_to_translate)[0]['translation_text']
        return jsonify({'translation': result})
    except Exception as e:
        return jsonify({'error': str(e)})

if __name__ == '__main__':
    app.run(debug=True)
